#!/usr/bin/env python3

import common
import lkmc.import_path
import shell_helpers
from shell_helpers import LF

class Main(common.TestCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Test and benchmark the Linux kernel boot. Use inits that exit immediately.
'''
        )
        self.add_argument(
            '--size',
            default=1,
            type=int,
            help='''\
See ./test --help for --size.
'''
        )

    def _bench(self, **kwargs):
        words = []
        for line in self.run.get_cli(**kwargs):
            words.extend(line)
        extra_params = shell_helpers.ShellHelpers().cmd_to_string(words + [LF])
        run_args = kwargs.copy()
        run_args.update(self.common_args)
        self.run_test(self.run, run_args, extra_params)

    def timed_main(self):
        # TODO bring this benchmark code back to life. Likely should go inside run with an option
        #gem5_insts() (
        #  printf "instructions $(./gem5-stat --arch "$1" sim_insts)\n" >> "$self.env['test_boot_benchmark_file']"
        #  newline
        #)
        #
        #qemu_insts() (
        #  common_arch="$1"
        #  ./qemu-trace2txt --arch "$common_arch"
        #  common_qemu_trace_txt_file="$("$getvar" --arch "$common_arch" qemu_trace_txt_file)"
        #  printf "instructions $(wc -l "${common_qemu_trace_txt_file}" | cut -d' ' -f1)\n" >> "$self.env['test_boot_benchmark_file']"
        #  newline
        #)
        #
        #rm -f "${self.env['test_boot_benchmark_file']}"
        self.run = lkmc.import_path.import_path_main('run')
        self.common_args = self.get_common_args()
        self.common_args['ctrl_c_host'] = True
        self.common_args['quit_after_boot'] = True
        if (self.env['emulator'] == 'qemu' or
                (self.env['emulator'] == 'gem5' and self.env['size'] >= 2)):
            self._bench()
        if self.env['host_arch'] == self.env['arch']:
            # TODO: find out why it fails.
            if self.env['emulator'] != 'gem5':
                self._bench(kvm=True)
        if self.env['emulator'] == 'qemu' and self.env['size'] >= 2:
            self._bench(trace='exec_tb')
        if self.env['emulator'] == 'gem5' and self.env['size'] >= 3:
            if self.env['arch'] == 'x86_64':
                cpu_types = [
                    # TODO segfault
                    #'DerivO3CPU'
                ]
            elif self.env['is_arm']:
                cpu_types = [
                    'DerivO3CPU',
                    'HPI',
                ]
            for cpu_type in cpu_types:
                self._bench(
                    extra_emulator_args=[
                        '--cpu-type', cpu_type,
                        '--caches',
                        '--l2cache',
                        '--l1d_size', '1024kB',
                        '--l1i_size', '1024kB',
                        '--l2_size', '1024kB',
                        '--l3_size', '1024kB',
                    ],
                )
            if self.env['arch'] == 'aarch64':
                # Do a fuller testing for aarch64.
                for build_type in ['debug', 'fast']:
                    self._bench(gem5_build_type=build_type)
                # Requires patching the executable.
                # self._bench(gem5_script='biglittle')

if __name__ == '__main__':
    Main().cli()
